package org.idea.qiyu.framework.shutdown.threadpool;

import org.idea.qiyu.framework.shutdown.config.IGracefulShutdown;

import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 具有优雅关闭特性的线程池
 *
 * @Author linhao
 * @Date created in 10:02 下午 2021/7/17
 */
public class GracefulThreadPoolExecutor implements IGracefulShutdown {

    private ThreadPoolExecutor threadPoolExecutor;

    private String threadPoolExecutorName;

    private static final AtomicBoolean isThreadPoolShutDown = new AtomicBoolean(false);

    /**
     * 默认关闭线程池超时位10秒
     */
    private long shutDownTimeOut = 10000;

    public GracefulThreadPoolExecutor(ThreadPoolExecutor threadPoolExecutor, String threadPoolExecutorName) {
        this.threadPoolExecutor = threadPoolExecutor;
        this.threadPoolExecutorName = threadPoolExecutorName;
    }

    public GracefulThreadPoolExecutor(ThreadPoolExecutor threadPoolExecutor, String threadPoolExecutorName, long shutDownTimeOut) {
        this.threadPoolExecutor = threadPoolExecutor;
        this.threadPoolExecutorName = threadPoolExecutorName;
        this.shutDownTimeOut = shutDownTimeOut;
    }

    public ThreadPoolExecutor getThreadPoolExecutor() {
        return threadPoolExecutor;
    }

    public void setThreadPoolExecutor(ThreadPoolExecutor threadPoolExecutor) {
        this.threadPoolExecutor = threadPoolExecutor;
    }

    @Override
    public String getName() {
        return threadPoolExecutorName;
    }

    @Override
    public void shutdown() {
        if (!isThreadPoolShutDown.compareAndSet(false, true)) {
            return;
        }
        threadPoolExecutor.shutdown();
        long beginCount = System.currentTimeMillis();
        while (threadPoolExecutor.getTaskCount()>=0 && (System.currentTimeMillis() - beginCount) <= shutDownTimeOut) {
            try {
                Thread.sleep(500);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }
}
